import matplotlib.pyplot as plt
import numpy as np
import offline_keypoint_detect_utils as utils
import os
import models.barcode_model_small as models
import torch

if __name__ == "__main__":
    postprocess_params = utils.PostProcessParams(
        width=128,
        height=128,
        heatValueThreshold=0.5,
        distThreshold=10.0,
        maxDistanceDifference=30,
        angleDifferenceThreshold=round(np.pi / 4, 2),
    )

    model_path = r"C:\Users\Administrator\Desktop\barcode_train\runs_barcode_1d\runs4\sigma_2epoch_40.pt"
    model = models.GrayDemoNet2()
    device = torch.device("cpu")
    model.to(device)
    checkpoint = torch.load(model_path, map_location=device)
    if isinstance(checkpoint, dict):
        model.load_state_dict(checkpoint["model_state_dict"])
    else:
        model = checkpoint

    img_directory_path = (
        r"C:\Users\Administrator\Desktop\barcode_train\real_barcode_data"
    )
    save_directory_path = (
        r"C:\Users\Administrator\Desktop\barcode_train\test_results\result19"
    )
    os.makedirs(save_directory_path, exist_ok=False)  # 防止测试结果覆盖

    images_name_list = os.listdir(img_directory_path)
    for img_name in images_name_list:
        if img_name.split(".")[-1] not in ["jpg", "png"]:
            continue
        print("Processing image: ", img_name)
        img_path = os.path.join(img_directory_path, img_name)
        img, outputs_np = utils.detect_image(
            img_path, model, postprocess_params, 128, 128
        )
        plt.subplot(2, 2, 1)
        plt.title("heatmap")
        plt.imshow(outputs_np[0, 0, :, :], cmap="gray")
        plt.subplot(2, 2, 2)
        plt.title("sin")
        plt.imshow(outputs_np[0, 1, :, :].squeeze(), cmap="gray")
        plt.subplot(2, 2, 3)
        plt.title("cos")
        plt.imshow(outputs_np[0, 2, :, :].squeeze(), cmap="gray")
        plt.subplot(2, 2, 4)
        plt.title("Visualization")
        plt.imshow(img, cmap="gray")
        plt.savefig(
            os.path.join(
                save_directory_path,
                img_name.split(".")[0]
                + "_"
                + str(postprocess_params.heatValueThreshold)
                + ".png",
            ),
            dpi=230,
        )
